import math
import torch
from torch import FloatTensor


XYZ2sRGB = FloatTensor([[ 3.2406, -1.5372, -0.4986],
                        [-0.9689,  1.8758,  0.0415],
                        [ 0.0557, -0.2040,  1.0570]])

# http://brucelindbloom.com/index.html?Eqn_RGB_XYZ_Matrix.html
ProPhotoRGB2XYZ = FloatTensor([[0.7976749, 0.1351917, 0.0313534],
                               [0.2880402, 0.7118741, 0.0000857],
                               [0.0000000, 0.0000000, 0.8252100]])

RGB2YUV = FloatTensor([[0.29900, 0.5870, 0.1140],
                       [-.33750, -.6625, 1.0000],
                       [1.00000, -.8374, -.1626]])

YUV2RGB = FloatTensor([[1.0, 0.0000, 0.7010],
                       [1.0, -.1721, -.3571],
                       [1.0, 0.8860, 0.0]])

xyz_color_matching = {
    "lambda": FloatTensor([390,395,400,405,410,415,420,425,430,435,440,445,450,455,460,465,470,475,480,485,490,495,500,505,510,515,520,525,530,535,540,545,550,555,560,565,570,575,580,585,590,595,600,605,610,615,620,625,630,635,640,645,650,655,660,665,670,675,680,685,690,695,700,705,710,715,720,725,730,735,740,745,750,755,760,765,770,775,780,785,790,795,800,805,810,815,820,825,830]),
    "xyz": FloatTensor([[0.003769647,0.009382967,0.02214302,0.04742986,0.08953803,0.1446214,0.2035729,0.2488523,0.2918246,0.3227087,0.3482554,0.3418483,0.3224637,0.2826646,0.2485254,0.2219781,0.1806905,0.129192,0.08182895,0.04600865,0.02083981,0.007097731,0.002461588,0.003649178,0.01556989,0.04315171,0.07962917,0.1268468,0.1818026,0.2405015,0.3098117,0.3804244,0.4494206,0.5280233,0.6133784,0.7016774,0.796775,0.8853376,0.9638388,1.051011,1.109767,1.14362,1.151033,1.134757,1.083928,1.007344,0.9142877,0.8135565,0.6924717,0.575541,0.4731224,0.3844986,0.2997374,0.2277792,0.1707914,0.1263808,0.09224597,0.0663996,0.04710606,0.03292138,0.02262306,0.01575417,0.01096778,0.00760875,0.005214608,0.003569452,0.002464821,0.001703876,0.001186238,0.000826954,0.00057583,0.00040583,0.000285658,0.000202185,0.000143827,0.000102469,7.34755E-05,5.25987E-05,3.80611E-05,2.75822E-05,2.00412E-05,1.45879E-05,1.06814E-05,7.85752E-06,5.76828E-06,4.25917E-06,3.16777E-06,2.35872E-06,1.76247E-06],
                        [0.000414616,0.001059646,0.002452194,0.004971717,0.00907986,0.01429377,0.02027369,0.02612106,0.03319038,0.0415794,0.05033657,0.05743393,0.06472352,0.07238339,0.08514816,0.1060145,0.1298957,0.1535066,0.1788048,0.2064828,0.237916,0.285068,0.3483536,0.4277595,0.5204972,0.6206256,0.718089,0.7946448,0.8575799,0.9071347,0.9544675,0.9814106,0.9890228,0.9994608,0.9967737,0.9902549,0.9732611,0.9424569,0.8963613,0.8587203,0.8115868,0.7544785,0.6918553,0.6270066,0.5583746,0.489595,0.4229897,0.3609245,0.2980865,0.2416902,0.1943124,0.1547397,0.119312,0.08979594,0.06671045,0.04899699,0.03559982,0.02554223,0.01807939,0.01261573,0.008661284,0.006027677,0.004195941,0.002910864,0.001995557,0.001367022,0.000944727,0.000653705,0.000455597,0.000317974,0.000221745,0.000156557,0.000110393,7.82744E-05,5.57886E-05,3.98188E-05,2.86018E-05,2.05126E-05,1.48724E-05,1.08E-05,7.86392E-06,5.73694E-06,4.2116E-06,3.10656E-06,2.28679E-06,1.69315E-06,1.26256E-06,9.42251E-07,7.05386E-07],
                        [0.0184726,0.04609784,0.109609,0.2369246,0.4508369,0.7378822,1.051821,1.305008,1.552826,1.74828,1.917479,1.918437,1.848545,1.664439,1.522157,1.42844,1.25061,0.9991789,0.7552379,0.5617313,0.4099313,0.3105939,0.2376753,0.1720018,0.1176796,0.08283548,0.05650407,0.03751912,0.02438164,0.01566174,0.00984647,0.006131421,0.003790291,0.002327186,0.001432128,0.000882253,0.000545242,0.000338674,0.000211777,0.000133503,8.49447E-05,5.46071E-05,3.54966E-05,2.33474E-05,1.55463E-05,1.04839E-05,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]])
}


# default 50% quality
default_jpeg_quantization_matrix = \
            FloatTensor([[16, 11, 10, 16, 24,  40,  51,  61],
                         [12, 12, 14, 19, 26,  58,  60,  55],
                         [14, 13, 16, 24, 40,  57,  69,  56],
                         [14, 17, 22, 29, 51,  87,  80,  62],
                         [18, 22, 37, 56, 68,  109, 103, 77],
                         [24, 35, 55, 64, 81,  104, 113, 92],
                         [49, 64, 78, 87, 103, 121, 120, 101],
                         [72, 92, 95, 98, 112, 100, 103, 99]])

# Photoshop quantization.
# https://www.impulseadventure.com/photo/jpeg-quantization.html
photoshop_jpeg_quantization_lum = \
    [
        # Luminance Level 0
        FloatTensor([
            [32, 33, 51, 81, 66, 39, 34, 17],
            [33, 36, 48, 47, 28, 23, 12, 12],
            [51, 48, 47, 28, 23, 12, 12, 12],
            [81, 47, 28, 23, 12, 12, 12, 12],
            [66, 28, 23, 12, 12, 12, 12, 12],
            [39, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 1
        FloatTensor([
            [27, 26, 41, 65, 66, 39, 34, 17],
            [26, 29, 38, 47, 28, 23, 12, 12],
            [41, 38, 47, 28, 23, 12, 12, 12],
            [65, 47, 28, 23, 12, 12, 12, 12],
            [66, 28, 23, 12, 12, 12, 12, 12],
            [39, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 2
        FloatTensor([
            [20, 17, 26, 41, 51, 39, 34, 17],
            [17, 18, 24, 39, 28, 23, 12, 12],
            [26, 24, 32, 28, 23, 12, 12, 12],
            [41, 39, 28, 23, 12, 12, 12, 12],
            [51, 28, 23, 12, 12, 12, 12, 12],
            [39, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 3
        FloatTensor([
            [18, 14, 22, 35, 44, 39, 34, 17],
            [14, 16, 21, 34, 28, 23, 12, 12],
            [22, 21, 27, 28, 23, 12, 12, 12],
            [35, 34, 28, 23, 12, 12, 12, 12],
            [44, 28, 23, 12, 12, 12, 12, 12],
            [39, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 4
        FloatTensor([
            [16, 11, 17, 27, 34, 39, 34, 17],
            [11, 12, 16, 26, 28, 23, 12, 12],
            [17, 16, 21, 28, 23, 12, 12, 12],
            [27, 26, 28, 23, 12, 12, 12, 12],
            [34, 28, 23, 12, 12, 12, 12, 12],
            [39, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 5
        FloatTensor([
            [12, 8, 13, 21, 26, 32, 34, 17],
            [8, 9, 12, 20, 27, 23, 12, 12],
            [13, 12, 16, 26, 23, 12, 12, 12],
            [21, 20, 26, 23, 12, 12, 12, 12],
            [26, 27, 23, 12, 12, 12, 12, 12],
            [32, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 6
        FloatTensor([
            [8, 6, 9, 14, 17, 21, 28, 17],
            [6, 6, 8, 13, 18, 23, 12, 12],
            [9, 8, 11, 17, 23, 12, 12, 12],
            [14, 13, 17, 23, 12, 12, 12, 12],
            [17, 18, 23, 12, 12, 12, 12, 12],
            [21, 23, 12, 12, 12, 12, 12, 12],
            [28, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 7
        FloatTensor([
            [10, 7, 11, 18, 22, 27, 34, 17],
            [7, 8, 10, 17, 23, 23, 12, 12],
            [11, 10, 14, 22, 23, 12, 12, 12],
            [18, 17, 22, 23, 12, 12, 12, 12],
            [22, 23, 23, 12, 12, 12, 12, 12],
            [27, 23, 12, 12, 12, 12, 12, 12],
            [34, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 8
        FloatTensor([
            [6, 4, 7, 11, 14, 17, 22, 17],
            [4, 5, 6, 10, 14, 19, 12, 12],
            [7, 6, 8, 14, 19, 12, 12, 12],
            [11, 10, 14, 19, 12, 12, 12, 12],
            [14, 14, 19, 12, 12, 12, 12, 12],
            [17, 19, 12, 12, 12, 12, 12, 12],
            [22, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 9
        FloatTensor([
            [4, 3, 4, 7, 9, 11, 14, 17],
            [3, 3, 4, 7, 9, 12, 12, 12],
            [4, 4, 5, 9, 12, 12, 12, 12],
            [7, 7, 9, 12, 12, 12, 12, 12],
            [9, 9, 12, 12, 12, 12, 12, 12],
            [11, 12, 12, 12, 12, 12, 12, 12],
            [14, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 10
        FloatTensor([
            [2, 2, 3, 4, 5, 6, 8, 11],
            [2, 2, 2, 4, 5, 7, 9, 11],
            [3, 2, 3, 5, 7, 9, 11, 12],
            [4, 4, 5, 7, 9, 11, 12, 12],
            [5, 5, 7, 9, 11, 12, 12, 12],
            [6, 7, 9, 11, 12, 12, 12, 12],
            [8, 9, 11, 12, 12, 12, 12, 12],
            [11, 11, 12, 12, 12, 12, 12, 12],
        ]),
        # Luminance Level 11
        FloatTensor([
            [1, 1, 1, 2, 3, 3, 4, 5],
            [1, 1, 1, 2, 3, 4, 4, 6],
            [1, 1, 2, 3, 4, 4, 5, 7],
            [2, 2, 3, 4, 4, 5, 7, 8],
            [3, 3, 4, 4, 5, 7, 8, 8],
            [3, 4, 4, 5, 7, 8, 8, 8],
            [4, 4, 5, 7, 8, 8, 8, 8],
            [5, 6, 7, 8, 8, 8, 8, 8],
        ]),
        # Luminance Level 12
        FloatTensor([
            [1, 1, 1, 1, 1, 1, 1, 2],
            [1, 1, 1, 1, 1, 1, 1, 2],
            [1, 1, 1, 1, 1, 1, 2, 2],
            [1, 1, 1, 1, 1, 2, 2, 3],
            [1, 1, 1, 1, 2, 2, 3, 3],
            [1, 1, 1, 2, 2, 3, 3, 3],
            [1, 1, 2, 2, 3, 3, 3, 3],
            [2, 2, 2, 3, 3, 3, 3, 3],
        ]),
    ]

photoshop_jpeg_quantization_chrom = \
    [
        # Chrominance Level 0
        FloatTensor([
            [34, 51, 52, 34, 20, 20, 17, 17],
            [51, 38, 24, 14, 14, 12, 12, 12],
            [52, 24, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 1
        FloatTensor([
            [29, 41, 52, 34, 20, 20, 17, 17],
            [41, 38, 24, 14, 14, 12, 12, 12],
            [52, 24, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 2
        FloatTensor([
            [21, 26, 33, 34, 20, 20, 17, 17],
            [26, 29, 24, 14, 14, 12, 12, 12],
            [33, 24, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 3
        FloatTensor([
            [20, 22, 29, 34, 20, 20, 17, 17],
            [22, 25, 24, 14, 14, 12, 12, 12],
            [29, 24, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 4
        FloatTensor([
            [17, 17, 22, 34, 20, 20, 17, 17],
            [17, 19, 22, 14, 14, 12, 12, 12],
            [22, 22, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 5
        FloatTensor([
            [13, 13, 17, 27, 20, 20, 17, 17],
            [13, 14, 17, 14, 14, 12, 12, 12],
            [17, 17, 14, 14, 12, 12, 12, 12],
            [27, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 6
        FloatTensor([
            [9, 9, 11, 18, 20, 20, 17, 17],
            [9, 10, 11, 14, 14, 12, 12, 12],
            [11, 11, 14, 14, 12, 12, 12, 12],
            [18, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 7
        FloatTensor([
            [11, 14, 31, 34, 20, 20, 17, 17],
            [14, 19, 24, 14, 14, 12, 12, 12],
            [31, 24, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 8
        FloatTensor([
            [7, 9, 19, 34, 20, 20, 17, 17],
            [9, 12, 19, 14, 14, 12, 12, 12],
            [19, 19, 14, 14, 12, 12, 12, 12],
            [34, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 9
        FloatTensor([
            [4, 6, 12, 22, 20, 20, 17, 17],
            [6, 8, 12, 14, 14, 12, 12, 12],
            [12, 12, 14, 14, 12, 12, 12, 12],
            [22, 14, 14, 12, 12, 12, 12, 12],
            [20, 14, 12, 12, 12, 12, 12, 12],
            [20, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
            [17, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 10
        FloatTensor([
            [3, 3, 7, 13, 15, 15, 15, 15],
            [3, 4, 7, 13, 14, 12, 12, 12],
            [7, 7, 13, 14, 12, 12, 12, 12],
            [13, 13, 14, 12, 12, 12, 12, 12],
            [15, 14, 12, 12, 12, 12, 12, 12],
            [15, 12, 12, 12, 12, 12, 12, 12],
            [15, 12, 12, 12, 12, 12, 12, 12],
            [15, 12, 12, 12, 12, 12, 12, 12],
        ]),
        # Chrominance Level 11
        FloatTensor([
            [1, 2, 4, 7, 8, 8, 8, 8],
            [2, 2, 4, 7, 8, 8, 8, 8],
            [4, 4, 7, 8, 8, 8, 8, 8],
            [7, 7, 8, 8, 8, 8, 8, 8],
            [8, 8, 8, 8, 8, 8, 8, 8],
            [8, 8, 8, 8, 8, 8, 8, 8],
            [8, 8, 8, 8, 8, 8, 8, 8],
            [8, 8, 8, 8, 8, 8, 8, 8],
        ]),
        # Chrominance Level 12
        FloatTensor([
            [1, 1, 1, 2, 3, 3, 3, 3],
            [1, 1, 1, 2, 3, 3, 3, 3],
            [1, 1, 2, 3, 3, 3, 3, 3],
            [2, 2, 3, 3, 3, 3, 3, 3],
            [3, 3, 3, 3, 3, 3, 3, 3],
            [3, 3, 3, 3, 3, 3, 3, 3],
            [3, 3, 3, 3, 3, 3, 3, 3],
            [3, 3, 3, 3, 3, 3, 3, 3],
        ]),
]

# 0-6 have subsampling, 7-12 don't.
photoshop_chroma_subsampling = [True] * 7 + [False] * 6



# DCT Coefficients
# The inverse coefficient is the same.
def _DCT_coeff():
    v = torch.arange(8).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).expand((8, 8, 8, 8)).float()
    u = torch.arange(8).unsqueeze( 0).unsqueeze(-1).unsqueeze(-1).expand((8, 8, 8, 8)).float()
    y = torch.arange(8).unsqueeze( 0).unsqueeze( 0).unsqueeze(-1).expand((8, 8, 8, 8)).float()
    x = torch.arange(8).unsqueeze( 0).unsqueeze( 0).unsqueeze( 0).expand((8, 8, 8, 8)).float()
    au = torch.ones((8, 8, 8, 8)).float()
    av = torch.ones((8, 8, 8, 8)).float()
    av[0, :, ...] = 0.707 # 1 / sqrt(2)
    au[:, 0, ...] = 0.707 # 1 / sqrt(2)

    coeff = au * av * torch.cos((2*x + 1)*u*math.pi/16.0) \
                    * torch.cos((2*y + 1)*v*math.pi/16.0)
    return coeff * 0.25

DCT_coeff = _DCT_coeff()
